# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: broker\worker_manager\on_premise_multi_slave.py
import logging, os, os.path, tempfile
from random import shuffle, random
from scanners.scan_app.errors import *
from broker.worker_manager.base import WorkerManager
from helpers.licensing import get_bxss_credentials
from scanners.scan_app.scan_app_base import *
from scanners.wvs.wvs_scan_app import WVSScanningApp
from broker.worker_manager.jobs_storage import JobsStorage
from broker.worker_manager.workers_storage import WorkersStorage
from settings_provider import settings
from db.tables.workers import LocalWorker
from scanners.constants import text_data_identifier
from broker.worker_manager.exceptions import *
from helpers.constants.scans import *
from db.tables.workers import WorkerStatus
from helpers.constants.jobs import ScanAppTypes

_logger = logging.getLogger("broker.op_worker_manager")


class WmOnPremiseStaticMultiSlave(WorkerManager):
    def has_slots(self, scanning_app):
        return True

    def __init__(
        self, shard_name, status_api_url, worker_storage, jobs_storage, debug_mode=False
    ):
        super().__init__(
            shard_name,
            status_api_url,
            worker_storage,
            jobs_storage,
            debug_mode=debug_mode,
        )
        self.local_wvs_scanning_app = WVSScanningApp(
            working_directory=settings.get(
                "wvs.temp_dir", default=tempfile.gettempdir()
            ),
            wvs_directory=settings.get(
                "wvs.app_dir",
                default=os.path.normpath(
                    os.getcwd() + os.path.sep + ".." + os.path.sep
                ),
            ),
            scan_storage_directory=settings.get(
                "storage.scans", default=tempfile.gettempdir()
            ),
        )
        self.supported_scanning_apps = {ScanAppTypes.WVS}

    def shutdown(self):
        self.local_wvs_scanning_app.shutdown()
        self.process_queued_jobs()

    def check_single_job_status(self, scan_session_id, scanning_app):
        """
        :param scanning_app:
        :param scan_session_id:
        :return: status
        """
        _logger.debug("check_job %s", scan_session_id)
        api_client = self.workers_storage.get_scan_api_client_by_job(
            self.jobs_storage, scan_session_id, scanning_app
        )
        if api_client is None:
            if scanning_app != ScanAppTypes.WVS:
                raise JobNotFound()
            response = self.local_wvs_scanning_app.get_job_status(
                ClsJobIdentifierList(job_ids=(scan_session_id,))
            )
        else:
            response = api_client.get_job_status(
                ClsJobIdentifierList(job_ids=(scan_session_id,))
            )
        for status in response.statuses:
            if status.job_id == scan_session_id:
                if status.status == "not_found":
                    raise JobNotFound()
                return status

        raise JobNotFound()

    def add_job(
        self,
        scan_session_id,
        scan_key,
        allowed_targets,
        target_id,
        jobs,
        scanning_app,
        scan_id,
        **kwargs
    ):
        """
        :param scan_id:
        :param scan_session_id:
        :param scan_key:
        :param allowed_targets:
        :param target_id:
        :param jobs:
        :param scanning_app:
        :param kwargs:
        :return:
        """
        if scanning_app not in self.supported_scanning_apps:
            raise WmScanningAppNotSupported("scanning app not supported")
        _logger.debug("add_job %s", scan_session_id)
        api_key = "0" + scan_key.replace("-", "")
        debug_mode_from_targets = False
        additional_targets = []
        if allowed_targets:
            for th, ti in allowed_targets.items():
                target = self._prepare_target(ti)
                if not target:
                    _logger.debug("target %s:%s not found", ti, th)
                    continue
                    debug_mode_from_targets |= target.configuration.get("debug", False)
                    additional_targets.append(target)

        main_target = self._prepare_target(target_id)
        if not main_target:
            _logger.error("target %s not found", target_id)
            raise WmTargetNotFound("target not found")
        debug_mode_from_targets |= main_target.configuration.get("debug", False)
        user_interaction = False
        if (
            kwargs.get("schedule_start_date") is None
            and kwargs.get("recurrence") is None
        ):
            manual_browsing_storage = settings.get("manual_browsing_storage")
            if manual_browsing_storage:
                ui_session_id = manual_browsing_storage.get_ui_session_id(scan_id)
                if ui_session_id:
                    user_interaction = True
        if jobs[scanning_app]["profile"] == "custom":
            scan_profile = text_data_identifier + ("\r\n").join(
                jobs[scanning_app]["checks"]
            )
        else:
            scan_profile = jobs[scanning_app]["profile"]
        job_description = dict(
            status_api_url=self.status_api_url + scan_session_id + "/wvs",
            status_api_key=api_key,
            job_id=scan_session_id,
            scan_profile=scan_profile,
            main_target=main_target,
            additional_targets=additional_targets,
            debug=self.debug_mode | debug_mode_from_targets,
            bxss_user_id=get_bxss_credentials()[1],
            bxss_scan_id=str(kwargs.get("acumonitor_id")),
            user_interaction=user_interaction,
        )
        if kwargs.get("status") == ScanStatusTypes.RESUMING:
            job_description["initial_state"] = (
                self.status_api_url
                + "downloads/"
                + self.generate_dependency_descriptor(
                    self.DEPENDENCY_SCAN_STATUS_DB, scanning_app + "." + scan_session_id
                )
            )
        job_description = ClsJobDescription(**job_description)
        worker_id_list = []
        run_locally = False
        fixed_worker_id = self._get_target_worker_id(target_id, scanning_app)
        if user_interaction:
            run_locally = True
            if fixed_worker_id not in (LocalWorker.WORKER_ID, None):
                raise ManualInterventionOnRemoteMachine()
            else:
                worker_id_list = self.workers_storage.get_worker_ids_by_scanning_app(
                    scanning_app, expected_status=WorkerStatus.ONLINE
                )
                if not worker_id_list:
                    run_locally = True
                    if fixed_worker_id not in (LocalWorker.WORKER_ID, None):
                        raise FixedWorkerNotFound()
                if fixed_worker_id:
                    if fixed_worker_id in (LocalWorker.WORKER_ID, None):
                        run_locally = True
                    else:
                        if fixed_worker_id not in worker_id_list:
                            raise FixedWorkerNotFound()
                        else:
                            worker_id_list = [fixed_worker_id]
                else:
                    if random() < 1 / (len(worker_id_list) + 1):
                        run_locally = True
                if run_locally:
                    self.local_wvs_scanning_app.start_job(job_description)
                    return LocalWorker.WORKER_ID
                shuffle(worker_id_list)
                for worker_id in worker_id_list:
                    try:
                        api_client = self.workers_storage.get_scan_api_client(worker_id)
                        api_client.start_job(job_description)
                        return worker_id
                    except NoMoreWorkerSlots:
                        pass
                    except WorkerNotAuthorized:
                        if fixed_worker_id:
                            if worker_id == fixed_worker_id:
                                raise FixedWorkerIsDetachedOrNotAuthorized()
                    except Exception as e:
                        logging.exception(
                            "managing worker %s failed with %s", worker_id, e
                        )
                        if fixed_worker_id:
                            if worker_id == fixed_worker_id:
                                worker_status_storage = settings.get(
                                    "worker_status_storage"
                                )
                                if worker_status_storage:
                                    ws = worker_status_storage.get(worker_id)
                                    if ws:
                                        if ws.is_offline:
                                            raise FixedWorkerIsOffline()

            raise NoMoreWorkers()

    def delete_job(self, scan_session_id, scanning_app):
        """
        :param scan_session_id:
        :param scanning_app:
        :return:
        """
        _logger.debug("deleting %s", scan_session_id)
        try:
            api_client = self.workers_storage.get_scan_api_client_by_job(
                self.jobs_storage, scan_session_id, scanning_app
            )
        except (JobNotFound, WorkerNotFound):
            return

        job_identifier = ClsJobIdentifier(job_id=scan_session_id)
        if api_client is None:
            pass
        try:
            self.local_wvs_scanning_app.delete_job(job_identifier)
            return
        except (JobNotFound, WorkerNotFound):
            return

        try:
            return api_client.delete_job(job_identifier)
        except (JobNotFound, WorkerNotFound):
            return
        except Exception as e:
            _logger.exception(
                "failed to delete %s:%s with %s", scan_session_id, scanning_app, e
            )
            self.queue_job(api_client.delete_job, 600, job_identifier)

    def abort_job(self, scan_session_id, scanning_app, reason=None):
        """
        :param scanning_app:
        :param scan_session_id:
        :param reason:
        :return: true if job was lost and retrying aborting does not make sense
        """
        _logger.debug("abort_job %s", scan_session_id)
        try:
            api_client = self.workers_storage.get_scan_api_client_by_job(
                self.jobs_storage, scan_session_id, scanning_app
            )
        except (JobNotFound, WorkerNotFound):
            return True

        job_identifier = ClsJobIdentifier(job_id=scan_session_id)
        if api_client is None:
            try:
                self.local_wvs_scanning_app.abort_job(job_identifier)
                return True
            except (JobNotFound, WorkerNotFound):
                return True

            try:
                api_client.abort_job(job_identifier)
            except (JobNotFound, WorkerNotFound):
                return True
            except Exception as e:
                _logger.exception(
                    "failed to abort %s:%s with %s", scan_session_id, scanning_app, e
                )
                self.queue_job(api_client.abort_job, 600, job_identifier)

            return False

    def pause_job(self, scan_session_id, scanning_app, reason=None):
        """
        :param scanning_app:
        :param scan_session_id:
        :param reason:
        :return:
        """
        _logger.debug("pause_job %s", scan_session_id)
        api_client = self.workers_storage.get_scan_api_client_by_job(
            self.jobs_storage, scan_session_id, scanning_app
        )
        job_identifier = ClsJobIdentifier(job_id=scan_session_id)
        if api_client is None:
            self.local_wvs_scanning_app.pause_job(job_identifier)
            return
        try:
            api_client.pause_job(job_identifier)
        except (JobNotFound, WorkerNotFound):
            raise
        except Exception as e:
            _logger.exception(
                "failed to pause %s:%s with %s", scan_session_id, scanning_app, e
            )
            self.queue_job(api_client.pause_job, 600, job_identifier)

    def download_scan_state_db_from_worker(self, scan_session_job):
        if scan_session_job["worker_id"] in (None, LocalWorker.WORKER_ID):
            return
        scan_state_storage = settings.get("storage.scan_state_db")
        if not scan_state_storage:
            logging.warning(
                "scan_state_storage is not defined, aborting scan state download"
            )
            return
        destination_path = os.path.join(
            scan_state_storage, scan_session_job["scanning_app"]
        )
        if not os.path.exists(destination_path):
            os.makedirs(destination_path, exist_ok=True)
        api_client = self.workers_storage.get_scan_api_client(
            scan_session_job["worker_id"]
        )
        if api_client:
            pass
        try:
            api_client.file_download(
                "/%s/state_db" % (scan_session_job["scan_session_id"],),
                filename=os.path.join(
                    destination_path, scan_session_job["scan_session_id"]
                ),
            )
        except Exception as e:
            _logger.exception("handle_remote_acumonitor_file failed with %s", e)
